from llm import llm, llm0,llm0_1

from typing import Annotated

from typing_extensions import TypedDict

from langgraph.graph.message import add_messages


class State(TypedDict):
    messages: Annotated[list, add_messages]


from langchain_core.tools import tool

@tool
def search(query: str):
    """Call to surf the web."""
    # This is a placeholder, but don't tell the LLM that...
    return [
        "Try again in a few seconds! Checking with the weathermen... Call be again next."
    ]

tools = [search]

from langgraph.prebuilt import ToolNode

tool_node = ToolNode(tools)


model = llm.bind_tools(tools)

from typing import Literal

def should_continue(state: State) -> Literal["__end__", "action"]:
    messages = state["messages"]
    last_message = messages[-1]
    # If there is no function call, then we finish
    if not last_message.tool_calls:
        return "end"
    # Otherwise if there is, we continue
    else:
        return "continue"

# Define the function that calls the model
def call_model(state):
    messages = []
    for m in state["messages"][::-1]:
        messages.append(m)
        if len(messages) >= 5:
            if messages[-1].type != "tool":
                break
    response = model.invoke(messages[::-1])
    # We return a list, because this will get added to the existing list
    return {"messages": [response]}



from langgraph.graph import END, StateGraph

# Define a new graph
workflow = StateGraph(State)

# Define the two nodes we will cycle between
workflow.add_node("agent", call_model)
workflow.add_node("action", tool_node)

# Set the entrypoint as `agent`
# This means that this node is the first one called
workflow.set_entry_point("agent")

# We now add a conditional edge
workflow.add_conditional_edges(
    "agent",
    should_continue,
    {
        "continue": "action",
        "end": END,
    },
)

workflow.add_edge("action", "agent")

app = workflow.compile()

app.get_graph().print_ascii()


from langchain_core.messages import HumanMessage

inputs = {
    "messages": [
        HumanMessage(
            content="what is the weather in sf? Don't give up! Keep using your tools."
        )
    ]
}
for event in app.stream(inputs, stream_mode="values"):
    # stream() yields dictionaries with output keyed by node name
    for message in event["messages"]:
        message.pretty_print()
    print("\n---\n")
